{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.cm import get_cmap\n",
    "import matplotlib.gridspec as gridspec\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, smooth2d, get_cartopy, cartopy_xlim, cartopy_ylim\n",
    "import glob\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import cartopy\n",
    "import cartopy.crs as ccrs\n",
    "import cmocean\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "from matplotlib.font_manager import FontProperties\n",
    "from mpl_toolkits.axes_grid.inset_locator import inset_axes\n",
    "from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes\n",
    "from mpl_toolkits.axes_grid1.inset_locator import mark_inset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = \"mynn\"\n",
    "dom = \"d01\"\n",
    "dom2 = \"d02\"\n",
    "sens = \"no_fractional_seaice\"\n",
    "day = \"28\"\n",
    "root = \"/glade/campaign/uwyo/wyom0122/lam/new_runs/\"\n",
    "if day == \"28\":\n",
    "#\troot = \"/glade/scratch/tjuliano/doe_comble/new_lam_runs/WRF/WRF-ASR-CAO/test/\"\n",
    "\tcases = [\"032820_no_edmf\",\"032820_mynn_l3\",\"032820\",\"032820\"]\n",
    "\tplt_titles = [r\"$\\Delta$x=3km, level 2.5 ED\", \"$\\Delta$x=3km, level 3 ED\", \"$\\Delta$x=3km, level 2.5 EDMF\", \"$\\Delta$x=1km, level 2.5 EDMF\"]\n",
    "else:\n",
    "#\troot = \"/glade/scratch/tjuliano/doe_comble/new_lam_runs/WRF/WRF-ASR-CAO/test/\"\n",
    "\tcases = [\"031220_no_edmf\",\"031220_mynn_l3\",\"031220\",\"031220\"]\n",
    "\tplt_titles = [r\"$\\Delta$x=3km, level 2.5 ED\", \"$\\Delta$x=3km, level 3 ED\", \"$\\Delta$x=3km, level 2.5 EDMF\", \"$\\Delta$x=1km, level 2.5 EDMF\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lwprng = np.arange(0,1000)\n",
    "times = [9,11,13,19,25,31,37,43,49]\n",
    "time_str = ['28_1030z','28_1130z','28_1230z','28_1530z','28_1830z','28_2130z','29_0030z','29_0330z','29_0630z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "lev = 10\n",
    "skipx = 25\n",
    "skipy = 50\n",
    "scale = 50\n",
    "barbw = 0.04\n",
    "g = 9.81\n",
    "Rd = 287.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Dataset(root + cases[2] + \"/wrfinput_\" + dom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosalpha = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha = s1.variables['SINALPHA'][-1,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lat2 = getvar(s1,\"lat\",meta=False)\n",
    "lon2 = getvar(s1,\"lon\",meta=False)\n",
    "hgt = getvar(s1,\"HGT\",meta=False)\n",
    "print (root + cases[2] + \"/wrfinput_\" + dom)\n",
    "sst = s1.variables['SST'][-1,:,:]\n",
    "lat_d02 = lat2\n",
    "lon_d02 = lon2\n",
    "hgt_d02 = hgt\n",
    "ll_lon = lon2[0,0]\n",
    "ll_lat = lat2[0,0]+0.1\n",
    "ul_lon = lon2[-1,0]\n",
    "ul_lat = lat2[-1,0]\n",
    "ur_lon = lon2[-1,-1]-0.15\n",
    "ur_lat = lat2[-1,-1]-0.1\n",
    "lr_lon = lon2[0,-1]\n",
    "lr_lat = lat2[0,-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NEW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_lat = (ll_lat + ur_lat)/2.\n",
    "#ref_lon = -(abs(ll_lon) + abs(ur_lon))/2.\n",
    "ref_lon = 8.5\n",
    "extent = [ll_lon,ur_lon,ll_lat,ur_lat]\n",
    "######################################################################################\n",
    "files_wind = sorted(glob.glob(root + cases[2] + \"/RESTART_E/wrfout_wind_\" + dom + \"*\"))\n",
    "files_cloud = sorted(glob.glob(root + cases[2] + \"/RESTART_E/wrfout_cloud_\" + dom + \"*\"))\n",
    "\n",
    "for i in np.arange(1):\n",
    "\tprint ('Reading file...')\n",
    "\tprint (files_wind[i])\n",
    "\ts1 = Dataset(files_wind[i])\n",
    "\ts2 = Dataset(files_cloud[i])\n",
    "\ts3 = Dataset(root + cases[2] + \"/wrfinput_\" + dom)\n",
    "    \n",
    "\tlat = getvar(s3,\"lat\",meta=False)\n",
    "\tlon = getvar(s3,\"lon\",meta=False)\n",
    "\tseaice = getvar(s3,\"SEAICE\",meta=False)\n",
    "\tseaice2 = getvar(s3,\"SEAICE\")\n",
    "\n",
    "\tu = getvar(s1,\"ua\",units=\"m s-1\",meta=False)\n",
    "\tv = getvar(s1,\"va\",units=\"m s-1\",meta=False)\n",
    "\tz = getvar(s1,\"z\", units=\"m\",meta=False)\n",
    "\n",
    "\ttemp = s1.variables['T'][-1,:,:,:] + 300.\n",
    "\trho = s1.variables['RHO'][-1,:,:,:]\n",
    "\tpres = rho * Rd * temp / 100.  # mb\n",
    "\n",
    "\t# Create the figure\n",
    "\t# create map subplots\n",
    "\t# 1000 MB\n",
    "\tfig = plt.figure(figsize=(20,15))\n",
    "\t#ax = fig.add_axes([0.05,0.05,0.9,0.9])\n",
    "\n",
    "\tmy_cmap = cmocean.cm.deep\n",
    "\tmy_cmap.set_under('lightskyblue')\n",
    "\t#my_cmap2 = cmo.thermal\n",
    "\n",
    "\t# Interpolate 100 m winds\n",
    "\tprint ('Interpolating...')\n",
    "\tplev = 925.\n",
    "\tu1000 = interplevel(u, pres, plev)\n",
    "\tv1000 = interplevel(v, pres, plev)\n",
    "\ttheta1000 = interplevel(temp, pres, plev)\n",
    "\tgph1000 = interplevel(z, pres, plev)\n",
    "\tgph1000 = smooth2d(gph1000, 9)\n",
    "    \n",
    "\tuearth1000 = u1000*cosalpha - v1000*sinalpha\n",
    "\tvearth1000 = v1000*cosalpha + u1000*sinalpha\n",
    "    \n",
    "\tm = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='h')\n",
    "    \n",
    "\t#urot1000, vrot1000 = m.rotate_vector(uearth1000,vearth1000,lon_d02,lat_d02)\n",
    "    \n",
    "\tplev = 500.\n",
    "\tu850 = interplevel(u, pres, plev)\n",
    "\tv850 = interplevel(v, pres, plev)\n",
    "\ttheta850 = interplevel(temp, pres, plev)\n",
    "\tgph850 = interplevel(z, pres, plev)\n",
    "\tgph850 = smooth2d(gph850, 9)\n",
    "    \n",
    "\tuearth850 = u850*cosalpha - v850*sinalpha\n",
    "\tvearth850 = v850*cosalpha + u850*sinalpha\n",
    "    \n",
    "\t#urot850, vrot850 = m.rotate_vector(uearth850,vearth850,lon_d02,lat_d02)\n",
    "    \n",
    "\ttheta850_mcao = interplevel(temp, pres, 850.)\n",
    "\tmcao = sst - theta850_mcao\n",
    "\tmcao[sst==0] = np.nan\n",
    "\tmcao[seaice>0] = np.nan\n",
    "    \n",
    "\t################################################################################################################\n",
    "\t############################################### LWP ############################################################\n",
    "\t################################################################################################################\n",
    "\tprint ('Plotting LWP...')\n",
    "    \n",
    "    \n",
    "\t#x_d01, y_d01 = m(lon_d01,lat_d01)\n",
    "\tx_d02, y_d02 = m(lon_d02,lat_d02)\n",
    "\tax1 = plt.subplot(1, 2, 1)\n",
    "\t############### 1000 MB\n",
    "\t# Make the contour plot\n",
    "\tif day == \"28\":\n",
    "\t\trng = np.arange(0,10040,40)\n",
    "\t\trng2 = np.arange(254,278,2)\n",
    "\telse:\n",
    "\t\trng = np.arange(0,10040,40)\n",
    "\t\trng2 = np.arange(252,276,2)\n",
    "\t#temp_contours = ax1.contourf(to_np(lon),to_np(lat),to_np(theta1000),levels=rng2,zorder=2,cmap=cmocean.cm.thermal,transform=ccrs.PlateCarree(),extend='both')\n",
    "\ttemp_contours = m.contourf(x_d02,y_d02,to_np(theta1000),levels=rng2,zorder=2,cmap=cmocean.cm.balance,extend='both')\n",
    "\tcbar = plt.colorbar(temp_contours,fraction=0.040, pad=0.04, orientation='horizontal', ax=ax1)\n",
    "\tcbar.ax.tick_params(labelsize=18)\n",
    "\tcbar.ax.set_xlabel('Potential Temperature (K)', labelpad=15, size=20)\n",
    "\t#gph_contours = ax1.contour(to_np(lon),to_np(lat),to_np(gph1000),levels=rng,zorder=2,colors='k',linewidths=3.,transform=ccrs.PlateCarree())\n",
    "\tgph_contours = m.contour(x_d02,y_d02,to_np(gph1000),levels=rng,zorder=2,colors='white',linewidths=3.)\n",
    "\tax1.clabel(gph_contours,rng[::2],fmt='%d',fontsize=14)\n",
    "#\tq = ax1.quiver(to_np(lon[::skipy,::skipx]), to_np(lat[::skipy,::skipx]), to_np(uearth1000[::skipy, ::skipx]), to_np(vearth1000[::skipy, ::skipx]), transform=ccrs.PlateCarree(), linewidth=1.5, zorder=4, units='inches', scale=scale, width=barbw,color='limegreen')\n",
    "\tq = ax1.quiver(x_d02[::skipy,::skipx], y_d02[::skipy,::skipx],\n",
    "\t\tuearth1000[::skipy, ::skipx], vearth1000[::skipy, ::skipx],\n",
    "\t\tlinewidth=1.5, zorder=5, units='inches', scale=scale, width=barbw,\n",
    "\t\tedgecolor='black', color='limegreen')\n",
    "\tqk = ax1.quiverkey (q, 1.05, -0.2, 25, '25 m/s', labelpos='E',fontproperties={'weight': 'bold','size':18},zorder=6)\n",
    "\t#ax1.contour(to_np(lon),to_np(lat),to_np(seaice),levels=[0.0001,0.25,0.5,0.75,1.0],colors=['skyblue','skyblue','skyblue','skyblue'],linewidths=3.,zorder=4,transform=ccrs.PlateCarree())\n",
    "\tm.contour(x_d02,y_d02,to_np(seaice),levels=[1.0],colors=['skyblue'],linewidths=3.,zorder=4)\n",
    "    \n",
    "\tcs = m.contourf(x_d02,y_d02,to_np(smooth2d(mcao,7)),np.arange(8,20,1),colors='none',hatches='x',zorder=4)\n",
    "    \n",
    "\tx_anx,y_anx = m(15.874211,69.232286)\n",
    "\tm.scatter(x_anx,y_anx,marker='D',s=80,facecolor='yellow',edgecolor='k',zorder=8)\n",
    "    \n",
    "\tfor ii, collection in enumerate(cs.collections):\n",
    "\t\tcollection.set_edgecolor('magenta')\n",
    "\t# Doing this also colors in the box around each level\n",
    "\t# We can remove the colored line around the levels by setting the linewidth to 0\n",
    "\tfor collection in cs.collections:\n",
    "\t\tcollection.set_linewidth(0.)\n",
    "    \n",
    "\t#m.drawmapboundary(fill_color='aqua')\n",
    "\tm.drawcoastlines(linewidth=2.0)\n",
    "\t#m.drawstates(linewidth=2.0)\n",
    "\tm.drawcountries(linewidth=2.0)\n",
    "\t#m.drawrivers(linewidth=3.0,color='blue')\n",
    "\tm.fillcontinents(color='gray',lake_color='gray',zorder=6)\n",
    "\tm.drawparallels(np.arange(60.,94.,4.0),labels=[True,False,False,False],fontsize=28,linewidth=3.0,zorder=7)\n",
    "\tm.drawmeridians(np.arange(-40.,70.,10.0),labels=[False,False,False,True],fontsize=28,linewidth=3.0,zorder=7)\n",
    "    \n",
    "\tax1.set_title(files_wind[i][-23:] + ', 925 hPa',fontdict = {'fontsize' : 28})\n",
    "    \n",
    "    \n",
    "\t#plt.show()\n",
    "\t#sys.exit()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "\tax2 = plt.subplot(1, 2, 2)\n",
    "\t############### 850 MB\n",
    "        # Make the contour plot\n",
    "\tif day == \"28\":\n",
    "\t\trng = np.arange(0,10080,80)\n",
    "\t\trng2 = np.arange(286,316,2)\n",
    "\telse:\n",
    "\t\trng = np.arange(0,10080,80)\n",
    "\t\trng2 = np.arange(282,295,1)\n",
    "\t#temp_contours = ax2.contourf(to_np(lon),to_np(lat),to_np(theta850),levels=rng2,zorder=2,cmap=cmocean.cm.thermal,transform=ccrs.PlateCarree(),extend='both')\n",
    "\ttemp_contours = m.contourf(x_d02,y_d02,to_np(theta850),levels=rng2,zorder=2,cmap=cmocean.cm.balance,extend='both')\n",
    "\tcbar = plt.colorbar(temp_contours,fraction=0.040, pad=0.04, orientation='horizontal', ax=ax2)\n",
    "\tcbar.ax.tick_params(labelsize=18)\n",
    "\tcbar.ax.set_xlabel('Potential Temperature (K)', labelpad=15, size=20)\n",
    "\t#gph_contours = ax2.contour(to_np(lon),to_np(lat),to_np(gph850),levels=rng,zorder=2,colors='k',linewidths=3.,transform=ccrs.PlateCarree())\n",
    "\tgph_contours = m.contour(x_d02,y_d02,to_np(gph850),levels=rng,zorder=2,colors='white',linewidths=3.)\n",
    "\tax2.clabel(gph_contours,rng[::2],fmt='%d',fontsize=14)\n",
    "#\tq = ax2.quiver(to_np(lon[::skipy,::skipx]), to_np(lat[::skipy,::skipx]), to_np(uearth850[::skipy, ::skipx]), to_np(vearth850[::skipy, ::skipx]), transform=ccrs.PlateCarree(), linewidth=1.5, zorder=4, units='inches', scale=scale, width=barbw,color='limegreen')\n",
    "\tq = ax2.quiver(x_d02[::skipy,::skipx], y_d02[::skipy,::skipx],\n",
    "\t\tuearth850[::skipy, ::skipx], vearth850[::skipy, ::skipx],\n",
    "\t\tlinewidth=1.5, zorder=4, units='inches', scale=scale, width=barbw,\n",
    "\t\tedgecolor='black', color='limegreen')\n",
    "#\tqk = ax2.quiverkey (q, 0.90, 0.04, 20, '20 m/s', labelpos='E',fontproperties={'weight': 'bold','size':18},zorder=6)\n",
    "\tax2.contour(x_d02,y_d02,to_np(seaice),levels=[1.0],colors=['skyblue'],linewidths=3.,zorder=4)\n",
    "    \n",
    "\tcs = m.contourf(x_d02,y_d02,to_np(smooth2d(mcao,7)),np.arange(8,20,1),colors='none',hatches='x',zorder=4)\n",
    "    \n",
    "\tm.scatter(x_anx,y_anx,marker='D',s=80,facecolor='yellow',edgecolor='k',zorder=8)\n",
    "    \n",
    "\tfor ii, collection in enumerate(cs.collections):\n",
    "\t\tcollection.set_edgecolor('magenta')\n",
    "\t# Doing this also colors in the box around each level\n",
    "\t# We can remove the colored line around the levels by setting the linewidth to 0\n",
    "\tfor collection in cs.collections:\n",
    "\t\tcollection.set_linewidth(0.)\n",
    "\n",
    "\t#m.drawmapboundary(fill_color='aqua')\n",
    "\tm.drawcoastlines(linewidth=2.0)\n",
    "\t#m.drawstates(linewidth=2.0)\n",
    "\tm.drawcountries(linewidth=2.0)\n",
    "\t#m.drawrivers(linewidth=3.0,color='blue')\n",
    "\tm.fillcontinents(color='gray',lake_color='gray',zorder=6)\n",
    "\tm.drawparallels(np.arange(60.,94.,4.0),labels=[True,False,False,False],fontsize=28,linewidth=3.0,zorder=7)\n",
    "\tm.drawmeridians(np.arange(-40.,70.,10.0),labels=[False,False,False,True],fontsize=28,linewidth=3.0,zorder=7)\n",
    "    \n",
    "\tax2.set_title(files_wind[i][-23:] + ', 500 hPa',fontdict = {'fontsize' : 28})\n",
    "\tplt.savefig(\"./figs_new_nz136/mar\" + day + \"_\" + dom + \"_synoptic_\" + str(i+1) + \"_BASEMAP.png\",dpi=300,bbox_inches='tight')\n",
    "\tplt.close()\n",
    "\t#plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
